package main

import (
	"context"
	"fmt"
	"log"
	"strconv"
	"time"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

type sampleSeriesType struct {
	Time            time.Time `bson:"_id" json:"-"`
	TimeUnixNanoStr string    `bson:"-" json:"timeUnixNano"`
	Samples         []float64 `bson:"samples" json:"samples"`
}

var client *mongo.Client = nil
var db *mongo.Database = nil

func initDB() error {
	// Set client options
	clientOptions := options.Client().ApplyURI(fmt.Sprintf("mongodb://%s:%s@%s/%s", conf.User, conf.Password, conf.HostPort, conf.DB))

	var err error
	// Connect to MongoDB
	client, err = mongo.Connect(context.Background(), clientOptions)
	if err != nil {
		return err
	}

	// Check the connection
	err = client.Ping(context.Background(), nil)
	if err != nil {
		return err
	}

	db = client.Database(conf.DB)
	log.Println("connected to MongoDB")
	return nil
}

func closeDB() {
	client.Disconnect(context.Background())
	db = nil
	client = nil
}

func saveDacSamples(data []float64) error {
	_, err := db.Collection("allSamples").InsertOne(context.TODO(), bson.M{
		"time":    time.Now(),
		"samples": data,
	})
	return err
}

func saveCentralServerSamples(data float64) error {
	_, err := db.Collection("centralServerSamples").InsertOne(context.TODO(), bson.M{
		"time":   time.Now(),
		"sample": data,
	})
	return err
}

// retriveDacSamples 获取某个时间前后的targetCount个数的样本点数据
func retriveDacSamples(centerTime time.Time, targetCount int) ([]*sampleSeriesType, error) {
	startTime := centerTime.Add(-time.Second)
	endTime := centerTime.Add(time.Second)
	// pipeline 先筛选centerTime前后1秒的数据，按与centerTime之差排序，选取targetCount个
	pipeline := mongo.Pipeline{
		bson.D{{
			Key: "$match",
			Value: bson.D{{
				Key:   "time",
				Value: bson.D{{Key: "$gte", Value: startTime}, {Key: "$lte", Value: endTime}},
			}},
		}},
		bson.D{{Key: "$unwind", Value: bson.D{{
			Key:   "path",
			Value: "$samples",
		}, {
			Key:   "includeArrayIndex",
			Value: "index",
		}}}},
		bson.D{{
			Key: "$addFields",
			Value: bson.M{
				"timeDis": bson.M{
					"$abs": bson.M{
						"$subtract": bson.A{"$time", centerTime},
					},
				},
			},
		}},
		bson.D{{
			Key: "$sort",
			Value: bson.D{
				{Key: "timeDis", Value: 1},
				{Key: "index", Value: 1},
			},
		}},
		bson.D{{Key: "$limit", Value: targetCount}},
		bson.D{{
			Key: "$group",
			Value: bson.M{
				"_id":     "$time",
				"samples": bson.M{"$push": "$samples"},
			},
		}},
		bson.D{{
			Key:   "$sort",
			Value: bson.D{{Key: "_id", Value: 1}},
		}},
	}
	cur, err := db.Collection("allSamples").Aggregate(context.TODO(), pipeline)
	if err != nil {
		return nil, err
	}

	results := []*sampleSeriesType{}
	for cur.Next(context.TODO()) {
		var samples sampleSeriesType
		if err := cur.Decode(&samples); err != nil {
			return nil, err
		}
		results = append(results, &samples)
	}

	if err := cur.Err(); err != nil {
		return nil, err
	}
	if err := cur.Close(context.TODO()); err != nil {
		return nil, err
	}

	for i, r := range results {
		results[i].TimeUnixNanoStr = strconv.FormatInt(r.Time.UnixNano(), 10)
	}

	return results, nil
}

// retriveDacSamplesBefore 获取某个时间点前的targetCount个数的样本点数据
// alreadyGot 该时间点上，之前已经获取了多少个数据点（用于避免重复获取数据）
func retriveDacSamplesBefore(endTime time.Time, alreadyGot int, targetCount int) ([]*sampleSeriesType, error) {
	// 先获取该时间点上的剩下的数据
	type oneSampleType struct {
		Time time.Time `bson:"time"`
		Data float64   `bson:"samples"`
	}
	pipelineOnTheTime := mongo.Pipeline{
		bson.D{{
			Key: "$match",
			Value: bson.D{{
				Key:   "time",
				Value: endTime,
			}},
		}},
		bson.D{{Key: "$unwind", Value: bson.D{{
			Key:   "path",
			Value: "$samples",
		}, {
			Key:   "includeArrayIndex",
			Value: "index",
		}}}},
		bson.D{{
			Key: "$sort",
			Value: bson.D{
				{Key: "index", Value: 1},
			},
		}},
	}
	cur, err := db.Collection("allSamples").Aggregate(context.TODO(), pipelineOnTheTime)
	if err != nil {
		return nil, err
	}
	samplesOnTheTime := []*oneSampleType{}
	for cur.Next(context.TODO()) {
		var sample oneSampleType
		if err := cur.Decode(&sample); err != nil {
			return nil, err
		}
		samplesOnTheTime = append(samplesOnTheTime, &sample)
	}
	if err := cur.Err(); err != nil {
		return nil, err
	}
	if err := cur.Close(context.TODO()); err != nil {
		return nil, err
	}

	// 该时间点的数据中，只留下之前没获取过的
	if len(samplesOnTheTime) <= alreadyGot {
		samplesOnTheTime = []*oneSampleType{}
	} else {
		samplesOnTheTime = samplesOnTheTime[:len(samplesOnTheTime)-alreadyGot]
	}

	// 获取除了该时间点的其余数据
	targetCount -= len(samplesOnTheTime)
	pipeline := mongo.Pipeline{
		bson.D{{
			Key: "$match",
			Value: bson.D{{
				Key:   "time",
				Value: bson.D{{Key: "$lt", Value: endTime}},
			}},
		}},
		bson.D{{Key: "$unwind", Value: bson.D{{
			Key:   "path",
			Value: "$samples",
		}, {
			Key:   "includeArrayIndex",
			Value: "index",
		}}}},
		bson.D{{
			Key: "$sort",
			Value: bson.D{
				{Key: "time", Value: -1},
				{Key: "index", Value: -1},
			},
		}},
		bson.D{{Key: "$limit", Value: targetCount}},
		bson.D{{
			Key: "$sort",
			Value: bson.D{
				{Key: "time", Value: 1},
				{Key: "index", Value: 1},
			},
		}},
		bson.D{{
			Key: "$group",
			Value: bson.M{
				"_id":     "$time",
				"samples": bson.M{"$push": "$samples"},
			},
		}},
		bson.D{{
			Key:   "$sort",
			Value: bson.D{{Key: "_id", Value: 1}},
		}},
	}
	cur, err = db.Collection("allSamples").Aggregate(context.TODO(), pipeline)
	if err != nil {
		return nil, err
	}

	results := []*sampleSeriesType{}
	if len(samplesOnTheTime) > 0 {
		// 先将该时间点的数据存进results
		data := make([]float64, len(samplesOnTheTime))
		for i, s := range samplesOnTheTime {
			data[i] = s.Data
		}
		var sample = sampleSeriesType{
			Time:    samplesOnTheTime[0].Time,
			Samples: data,
		}
		results = append(results, &sample)
	}
	// 然后将其余时间的数据存进results
	for cur.Next(context.TODO()) {
		var samples sampleSeriesType
		if err := cur.Decode(&samples); err != nil {
			return nil, err
		}
		results = append(results, &samples)
	}

	if err := cur.Err(); err != nil {
		return nil, err
	}
	if err := cur.Close(context.TODO()); err != nil {
		return nil, err
	}

	for i, r := range results {
		results[i].TimeUnixNanoStr = strconv.FormatInt(r.Time.UnixNano(), 10)
	}

	return results, nil
}

// retriveDacSamplesAfter 获取某个时间点后的targetCount个数的样本点数据
// alreadyGot 该时间点上，之前已经获取了多少个数据点（用于避免重复获取数据）
func retriveDacSamplesAfter(startTime time.Time, alreadyGot int, targetCount int) ([]*sampleSeriesType, error) {
	// 先获取该时间点上的剩下的数据
	type oneSampleType struct {
		Time time.Time `bson:"time"`
		Data float64   `bson:"samples"`
	}
	pipelineOnTheTime := mongo.Pipeline{
		bson.D{{
			Key: "$match",
			Value: bson.D{{
				Key:   "time",
				Value: startTime,
			}},
		}},
		bson.D{{Key: "$unwind", Value: bson.D{{
			Key:   "path",
			Value: "$samples",
		}, {
			Key:   "includeArrayIndex",
			Value: "index",
		}}}},
		bson.D{{
			Key: "$sort",
			Value: bson.D{
				{Key: "index", Value: 1},
			},
		}},
	}
	cur, err := db.Collection("allSamples").Aggregate(context.TODO(), pipelineOnTheTime)
	if err != nil {
		return nil, err
	}
	samplesOnTheTime := []*oneSampleType{}
	for cur.Next(context.TODO()) {
		var sample oneSampleType
		if err := cur.Decode(&sample); err != nil {
			return nil, err
		}
		samplesOnTheTime = append(samplesOnTheTime, &sample)
	}
	if err := cur.Err(); err != nil {
		return nil, err
	}
	if err := cur.Close(context.TODO()); err != nil {
		return nil, err
	}

	// 该时间点的数据中，只留下之前没获取过的
	if len(samplesOnTheTime) <= alreadyGot {
		samplesOnTheTime = []*oneSampleType{}
	} else {
		samplesOnTheTime = samplesOnTheTime[alreadyGot:]
	}

	// 获取除了该时间点的其余数据
	targetCount -= len(samplesOnTheTime)
	pipeline := mongo.Pipeline{
		bson.D{{
			Key: "$match",
			Value: bson.D{{
				Key:   "time",
				Value: bson.D{{Key: "$gt", Value: startTime}},
			}},
		}},
		bson.D{{Key: "$unwind", Value: bson.D{{
			Key:   "path",
			Value: "$samples",
		}, {
			Key:   "includeArrayIndex",
			Value: "index",
		}}}},
		bson.D{{
			Key: "$sort",
			Value: bson.D{
				{Key: "time", Value: 1},
				{Key: "index", Value: 1},
			},
		}},
		bson.D{{Key: "$limit", Value: targetCount}},
		bson.D{{
			Key: "$group",
			Value: bson.M{
				"_id":     "$time",
				"samples": bson.M{"$push": "$samples"},
			},
		}},
		bson.D{{
			Key:   "$sort",
			Value: bson.D{{Key: "_id", Value: 1}},
		}},
	}
	cur, err = db.Collection("allSamples").Aggregate(context.TODO(), pipeline)
	if err != nil {
		return nil, err
	}

	results := []*sampleSeriesType{}
	if len(samplesOnTheTime) > 0 {
		// 先将该时间点的数据存进results
		data := make([]float64, len(samplesOnTheTime))
		for i, s := range samplesOnTheTime {
			data[i] = s.Data
		}
		var sample = sampleSeriesType{
			Time:    samplesOnTheTime[0].Time,
			Samples: data,
		}
		results = append(results, &sample)
	}
	// 然后将其余时间的数据存进results
	for cur.Next(context.TODO()) {
		var samples sampleSeriesType
		if err := cur.Decode(&samples); err != nil {
			return nil, err
		}
		results = append(results, &samples)
	}

	if err := cur.Err(); err != nil {
		return nil, err
	}
	if err := cur.Close(context.TODO()); err != nil {
		return nil, err
	}

	for i, r := range results {
		results[i].TimeUnixNanoStr = strconv.FormatInt(r.Time.UnixNano(), 10)
	}

	return results, nil
}
